﻿# -*- coding: utf-8 -*-   # 使用UTF-8编码以支持所有字符
import numpy as np
import matplotlib.pyplot as plt
from scipy import integrate
from typing import Callable, List, Tuple, Optional

class StationarySchroedingerEquation:
    """
    定态薛定谔方程类，用于处理和求解一维定态薛定谔方程
    
    属性:
    - mass: 粒子质量
    - hbar: 约化普朗克常数
    - potential: 势能函数 V(x)
    - x_range: 空间范围 (x_start, x_end)
    - x_points: 空间网格点数量
    """
    
    def __init__(self, mass: float = 1.0, hbar: float = 1.0):
        """
        初始化定态薛定谔方程类
        
        参数:
        - mass: 粒子质量，默认值为1.0
        - hbar: 约化普朗克常数，默认值为1.0
        """
        self.mass = mass
        self.hbar = hbar
        self.potential = None
        self.x_range = None
        self.x_points = None
    
    def set_potential(self, potential: Callable[[float], float]) -> None:
        """
        设置势能函数
        
        参数:
        - potential: 势能函数 V(x)，接受一个浮点数参数x，返回浮点数
        """
        self.potential = potential
    
    def set_space_grid(self, x_start: float, x_end: float, x_points: int) -> None:
        """
        设置空间网格
        
        参数:
        - x_start: 空间起始点
        - x_end: 空间结束点
        - x_points: 网格点数量
        """
        self.x_range = (x_start, x_end)
        self.x_points = x_points
    
    def _build_hamiltonian(self) -> np.ndarray:
        """
        构建哈密顿矩阵
        
        在量子力学中，哈密顿算符 H = T + V，其中：
        - T 是动能算符 T = -ħ²/(2m)·d²/dx²
        - V 是势能算符 V(x)
        
        在数值计算中，我们将连续的算符转换为离散的矩阵表示，这一过程称为有限差分方法。
        
        返回:
        - hamiltonian: 哈密顿矩阵
        """
        if self.potential is None or self.x_range is None or self.x_points is None:
            raise ValueError("请先设置势能函数和空间网格")
        
        # 获取空间范围和计算网格间距
        x_start, x_end = self.x_range
        dx = (x_end - x_start) / (self.x_points - 1)  # 网格点之间的间距
        x = np.linspace(x_start, x_end, self.x_points)  # 生成空间网格点
        
        # 构建动能项 (使用中心差分近似二阶导数)
        # 中心差分公式: d²ψ/dx² ≈ [ψ(x+dx) - 2ψ(x) + ψ(x-dx)]/dx²
        # 对应的矩阵表示为三对角矩阵：
        # 主对角线：-2/dx²
        # 上下对角线：1/dx²
        kinetic = - (self.hbar**2) / (2 * self.mass) * (
            np.diag(np.ones(self.x_points - 1), -1)  # 下对角线：全1
            - 2 * np.diag(np.ones(self.x_points))     # 主对角线：全-2
            + np.diag(np.ones(self.x_points - 1), 1)  # 上对角线：全1
        ) / dx**2  # 除以dx²完成差分近似
        
        # 构建势能项
        # 势能在离散表示中为对角矩阵，对角线元素为各网格点的势能值
        potential = np.diag(self.potential(x))  # 势能函数在各网格点的值构成对角矩阵
        
        # 总哈密顿量：动能矩阵 + 势能矩阵
        # 这对应于量子力学中的哈密顿算符 H = T + V
        hamiltonian = kinetic + potential
        
        return hamiltonian
    
    def solve_eigenvalues(self, n: int = 5) -> Tuple[np.ndarray, np.ndarray]:
        """
        求解前n个本征值和本征函数
        
        参数:
        - n: 要求解的本征值数量
        
        返回:
        - eigenvalues: 前n个本征值（能量）
        - eigenfunctions: 对应的本征函数
        """
        hamiltonian = self._build_hamiltonian()
        
        # 使用numpy的eigh函数求解厄米矩阵的本征值和本征函数
        # 这里对应量子力学中的定态薛定谔方程：Hψ = Eψ
        # 其中：
        # - H是哈密顿算符（在代码中表示为哈密顿矩阵）
        # - ψ是本征函数（波函数）
        # - E是本征值（能量）
        # np.linalg.eigh专门用于求解厄米矩阵的本征值和本征函数，
        # 它返回按升序排列的本征值和对应的归一化本征向量
        eigenvalues, eigenfunctions = np.linalg.eigh(hamiltonian)
        
        # 返回前n个本征值和本征函数
        return eigenvalues[:n], eigenfunctions[:, :n]
    
    def normalize_wavefunction(self, wavefunction: np.ndarray) -> np.ndarray:
        """
        归一化波函数
        
        在量子力学中，归一化常数N是确保波函数满足概率诠释的关键因子。根据玻恩定则，
        |ψ(x)|² 表示粒子在位置x处出现的概率密度。为了保证粒子在整个空间中出现的总概率为1，
        波函数必须满足归一化条件：
        
        ∫|ψ_normalized(x)|² dx = 1 （从负无穷到正无穷积分）
        
        归一化过程的数学表达式为：
        
        ψ_normalized(x) = ψ(x) / N
        
        其中归一化常数N由下式确定：
        
        N = √(∫|ψ(x)|² dx) ，从x_start到x_end积分
        
        归一化常数N在以下量子力学公式中至关重要：
        1. 概率计算：P(a ≤ x ≤ b) = ∫|ψ_normalized(x)|² dx （从a到b积分）
        2. 算符期望值：<O> = ∫ψ_normalized*(x) O ψ_normalized(x) dx
        3. 波函数正交性：∫ψ_m*(x) ψ_n(x) dx = δ_mn （克罗内克函数）
        
        参数:
        - wavefunction: 待归一化的波函数 ψ(x)
        
        返回:
        - normalized_wavefunction: 归一化后的波函数 ψ_normalized(x)
        """
        if self.x_range is None or self.x_points is None:
            raise ValueError("请先设置空间网格")
        
        x_start, x_end = self.x_range
        dx = (x_end - x_start) / (self.x_points - 1)
        
        # 计算归一化常数 N = √(∫|ψ(x)|² dx)
        # 注意：这行代码不是直接表示x点上的波函数值ψ(x)，而是计算归一化常数
        # 1. np.abs(wavefunction)**2 计算每个x点上波函数的模平方 |ψ(x)|²
        # 2. np.trapz(..., dx=dx) 使用梯形法则计算 |ψ(x)|² 在整个空间范围的积分 ∫|ψ(x)|² dx
        # 3. np.sqrt(...) 对积分结果取平方根，得到归一化常数 N
        norm = np.sqrt(np.trapz(np.abs(wavefunction)**2, dx=dx))
        
        # 返回归一化后的波函数 ψ_normalized(x) = ψ(x) / N
        return wavefunction / norm
    
    def plot_wavefunctions(self, eigenvalues: np.ndarray, eigenfunctions: np.ndarray, 
                          save_path: Optional[str] = None) -> None:
        """
        绘制本征函数和相应的能量水平
        
        参数:
        - eigenvalues: 本征值（能量）
        - eigenfunctions: 本征函数
        - save_path: 保存路径，若为None则不保存
        """
        if self.x_range is None or self.x_points is None or self.potential is None:
            raise ValueError("请先设置空间网格和势能函数")
        
        x_start, x_end = self.x_range
        x = np.linspace(x_start, x_end, self.x_points)
        v = self.potential(x)
        
        # 创建图形
        plt.figure(figsize=(12, 10))
        
        # 绘制势能
        plt.plot(x, v, 'k--', label='势能 V(x)')
        
        # 绘制各个本征函数（向上平移对应能量值）
        for i in range(len(eigenvalues)):
            # 归一化本征函数并缩放以便于显示
            psi = self.normalize_wavefunction(eigenfunctions[:, i])
            scale = (np.max(v) - np.min(v)) / 4  # 缩放因子
            plt.plot(x, eigenvalues[i] + psi * scale, label=f'波函数 ψ_{i}(x) (E={eigenvalues[i]:.3f})')
            
            # 绘制能量水平线
            plt.axhline(y=eigenvalues[i], color='gray', linestyle='-', alpha=0.3)
        
        # 设置图形属性
        plt.xlabel('位置 x')
        plt.ylabel('能量 / 波函数')
        plt.title('定态薛定谔方程的本征函数和能量水平')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 保存图形（如果提供了保存路径）
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        plt.show()
    
    def calculate_expectation_value(self, wavefunction: np.ndarray, operator: Callable[[np.ndarray], np.ndarray]) -> float:
        """
        计算算符的期望值
        
        参数:
        - wavefunction: 波函数
        - operator: 算符，接受一个波函数数组，返回作用后的结果
        
        返回:
        - expectation_value: 期望值
        """
        if self.x_range is None or self.x_points is None:
            raise ValueError("请先设置空间网格")
        
        x_start, x_end = self.x_range
        dx = (x_end - x_start) / (self.x_points - 1)
        
        # 归一化波函数
        normalized_psi = self.normalize_wavefunction(wavefunction)
        
        # 计算期望值 <ψ|O|ψ>
        expectation_value = np.trapz(np.conj(normalized_psi) * operator(normalized_psi), dx=dx)
        
        return expectation_value.real  # 期望值应为实数
        
    def get_time_dependent_factor(self, energy: float, time: float) -> np.ndarray:
        """
        计算波函数的时间相关部分
        
        在量子力学中，定态波函数的时间演化因子为：φ(t) = e^(-iEt/hbar)
        其中：
        - E 是能量本征值
        - t 是时间
        - hbar 是约化普朗克常数
        
        参数:
        - energy: 能量本征值
        - time: 时间值
        
        返回:
        - time_factor: 时间相关的相位因子
        """
        # 计算时间演化因子
        time_factor = np.exp(-1j * energy * time / self.hbar)
        
        return time_factor
    
    def get_spatial_wavefunction(self, eigenfunction: np.ndarray, normalize: bool = True) -> np.ndarray:
        """
        获取空间相关波函数，可以选择是否归一化
        
        空间波函数对应于定态薛定谔方程的本征函数 ψ(x)，它描述了量子态的空间分布
        
        参数:
        - eigenfunction: 从求解本征值问题得到的原始本征函数
        - normalize: 是否对波函数进行归一化（默认为True）
        
        返回:
        - spatial_wavefunction: 空间波函数（归一化或原始形式）
        
        注意:
        - 在量子力学中，波函数通常需要归一化以满足概率诠释（总概率为1）
        - 如果您已经有了归一化好的点集，可以设置normalize=False来直接使用
        """
        if normalize:
            # 调用已有的归一化方法处理空间波函数
            spatial_wavefunction = self.normalize_wavefunction(eigenfunction)
        else:
            # 直接返回原始波函数，不进行归一化
            spatial_wavefunction = np.copy(eigenfunction)
        
        return spatial_wavefunction
    
    def set_normalized_wavefunction(self, normalized_wavefunction: np.ndarray) -> np.ndarray:
        """
        直接设置已归一化的波函数
        
        此方法用于当用户已经有了归一化好的波函数点集时，直接使用这些点集
        而不需要再次归一化
        
        参数:
        - normalized_wavefunction: 已经归一化好的波函数点集
        
        返回:
        - validated_wavefunction: 经过验证的归一化波函数
        
        注意:
        - 此方法会验证输入波函数的归一化程度
        - 如果波函数偏离归一化状态较远，会发出警告
        """
        if self.x_range is None or self.x_points is None:
            raise ValueError("请先设置空间网格")
        
        x_start, x_end = self.x_range
        dx = (x_end - x_start) / (self.x_points - 1)
        
        # 验证波函数的归一化程度
        norm = np.sqrt(np.trapz(np.abs(normalized_wavefunction)**2, dx=dx))
        
        # 如果归一化程度偏离1较远，发出警告
        if abs(norm - 1.0) > 1e-3:
            import warnings
            warnings.warn(f"输入的波函数归一化因子为{norm}，偏离1较远，可能不是正确归一化的波函数")
        
        return normalized_wavefunction
    
    def construct_full_wavefunction(self, spatial_wavefunction: np.ndarray, energy: float, time: float) -> np.ndarray:
        """
        构建完整的量子波函数，结合空间和时间两部分
        
        量子力学中，完整的波函数可以表示为空间部分和时间部分的乘积：
        Ψ(x, t) = ψ(x) * φ(t)
        其中：
        - ψ(x) 是归一化的空间波函数
        - φ(t) 是时间演化因子
        
        参数:
        - spatial_wavefunction: 归一化的空间波函数
        - energy: 对应的能量本征值
        - time: 时间值
        
        返回:
        - full_wavefunction: 完整的波函数 Ψ(x, t)
        """
        # 获取时间相关因子
        time_factor = self.get_time_dependent_factor(energy, time)
        
        # 完整波函数是空间部分和时间部分的乘积
        full_wavefunction = spatial_wavefunction * time_factor
        
        return full_wavefunction
    
    def get_full_wavefunction(self, spatial_wavefunction: np.ndarray, energy: float, time: float) -> np.ndarray:
        """
        获取完整的波函数，包含时间和空间部分的乘积（为保持向后兼容性保留此方法）
        
        参数:
        - spatial_wavefunction: 空间部分的波函数（归一化的本征函数）
        - energy: 对应的能量本征值
        - time: 时间值
        
        返回:
        - full_wavefunction: 完整的波函数 Ψ(x, t)
        """
        return self.construct_full_wavefunction(spatial_wavefunction, energy, time)


# 提供一些常见的势能函数作为示例
def infinite_well_potential(x, width: float = 1.0, height: float = 100.0):
    """
    无限深势阱势能（支持标量和NumPy数组输入）
    
    参数:
    - x: 位置（标量或NumPy数组）
    - width: 势阱宽度
    - height: 势阱外的势能高度
    
    返回:
    - 势能值（与x相同类型）
    """
    half_width = width / 2
    
    # 处理NumPy数组输入
    if isinstance(x, np.ndarray):
        # 创建与x相同形状的结果数组
        result = np.zeros_like(x)
        # 找出满足条件的元素索引
        mask = (-half_width <= x) & (x <= half_width)
        # 根据条件设置势能值
        result[mask] = 0.0
        result[~mask] = height
        return result
    # 处理标量输入
    else:
        return 0.0 if -half_width <= x <= half_width else height


def harmonic_oscillator_potential(x: float, omega: float = 1.0) -> float:
    """
    简谐振子势能
    
    参数:
    - x: 位置
    - omega: 角频率
    
    返回:
    - 势能值 V(x) = 0.5 * m * omega^2 * x^2
    """
    return 0.5 * omega**2 * x**2


def square_barrier_potential(x: float, width: float = 0.5, height: float = 5.0, center: float = 0.0) -> float:
    """
    方形势垒势能
    
    参数:
    - x: 位置
    - width: 势垒宽度
    - height: 势垒高度
    - center: 势垒中心位置
    
    返回:
    - 势能值
    """
    half_width = width / 2
    return height if (center - half_width) <= x <= (center + half_width) else 0.0


# 注意：此类设计为作为库使用，不包含主程序入口
# 如需使用示例，可以在其他Python文件中导入并使用此类

# 示例导入方式：
# from quantum_mechanics.stationary_schroedinger import StationarySchroedingerEquation, harmonic_oscillator_potential
# 
# # 创建实例
# se = StationarySchroedingerEquation()
# # 设置势能函数
# se.set_potential(lambda x: harmonic_oscillator_potential(x))
# # 设置空间网格
# se.set_space_grid(-5, 5, 1000)
# # 求解并使用其他功能
# eigenvalues, eigenfunctions = se.solve_eigenvalues(n=5)